import pyrealsense2 as rs
import numpy as np
import cv2

# 获取 ultralytics 所在的父目录
ultralytics_path = "/ultralytics/ultralytics"  # 请替换为 ultralytics 的实际路径
# 确保路径存在
if os.path.exists(ultralytics_path):
    # 将 ultralytics 的父目录添加到 sys.path
    # sys.path.insert(0, ...) 会将路径添加到列表的开头，确保优先被查找
    sys.path.insert(0, ultralytics_path)
    print(f"Added ultralytics path to sys.path: {ultralytics_path}")
else:
    print(f"Warning: ultralytics path not found at {ultralytics_path}")
# 现在可以正常导入 ultralytics 了
from ultralytics import YOLO


# 创建管道
pipeline = rs.pipeline()

# 创建配置对象
config = rs.config()

# 启用彩色和深度流
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

# 开始流传输
profile = pipeline.start(config)

# 获取深度传感器的深度标尺（单位：米）
depth_sensor = profile.get_device().first_depth_sensor()
depth_scale = depth_sensor.get_depth_scale()
print(f"Depth Scale: {depth_scale}")
try:
    while True:
        # 等待一组帧（深度和彩色）
        frames = pipeline.wait_for_frames()
        
        # 获取深度帧和彩色帧
        depth_frame = frames.get_depth_frame()
        color_frame = frames.get_color_frame()
        
        if not depth_frame or not color_frame:
            continue
            
        # 转换为numpy数组
        depth_image = np.asanyarray(depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())
        
        # 应用颜色映射到深度图像（用于可视化）
        depth_colormap = cv2.applyColorMap(
            cv2.convertScaleAbs(depth_image, alpha=0.03), 
            cv2.COLORMAP_JET)
        
        # 显示图像
        cv2.imshow('Color', color_image)
        cv2.imshow('Depth', depth_colormap)
        
        if cv2.waitKey(1) == ord('q'):
            break
finally:
    # 停止管道
    pipeline.stop()
    cv2.destroyAllWindows()
